import sys, copy
from collections import defaultdict

def handle(alp, n):
    if alp in stack:
        del stack[n-1:]
        del classid[n-1:]
    if len(stack) >= n:
        if stack[n-1] == '-':
            del stack[n-1:]
            del classid[n-1:]
    stack.append(arr[3])
    classid.append(e)
    if len(stack) < n:
        num = n - len(stack)
        [stack.insert(len(stack)-1, '-') for n in range(num)]
        [classid.insert(len(classid)-1, '-') for n in range(num)]

outpre = sys.argv[1].split('.')[0]
en = ''
#rawseq = ['U', 'R', 'D', 'P', 'C', 'O', 'F', 'G', 'S']
#lenD, lenP, lenC, lenO, lenF, lenG, lenS = 3, 4, 5, 6, 7, 8, 9
tax2name = {}
with open(sys.argv[1]) as kraken:
    ks = defaultdict(int)
    for i in kraken:
        if not i.startswith('U'):
            item = i.split('\t')
            en = item[2].split(' (taxid')[0]
            t = item[2].split()[-1].replace(')', '')
            s = item[4].split()
            tsum = 0
            allsum = 0
            if en.find('Mycobacterium tuberculosis') > -1:
                for n in s:
                    st, kn = n.split(':')
                    if st == 'A':
                        continue
                    if st in ['77643', '1773']:
                        tsum += int(kn)
                    allsum += int(kn)
                if tsum / allsum > 0.5:
                    ks['Mycobacterium tuberculosis'] += 1
            elif en.find('Mycobacterium avium') > -1:
                for n in s:
                    st, kn = n.split(':')
                    if st == 'A':
                        continue
                    if st in ['120793', '1764']:
                        tsum += int(kn)
                    allsum += int(kn)
                if tsum / allsum > 0.5:
                    ks['Mycobacterium avium'] += 1
            else:
                for n in s:
                    st, kn = n.split(':')
                    if st == 'A':
                        continue
                    if st == t:
                        tsum += int(kn)
                    allsum += int(kn)
                if tsum / allsum > 0.5:
                    ks[en] += 1

sample = defaultdict(list)
with open(sys.argv[2]) as kreport:
    uc = kreport.readline().strip().split('\t')
    rc = kreport.readline().strip().split('\t')
    stack = [uc[3], rc[3]]
    classid = ['_'.join(uc[3:6]), '_'.join(rc[3:6])]
    for i in kreport:
        arr = i.strip().split('\t')
        ss = arr[5].strip()
        s = '_'.join(arr[5].strip().split(' '))
        e = f'{arr[3]}_{arr[4]}_{s}'
        if arr[3] == 'S':
            stack.append(arr[3])
            classid.append(e)
            if len(stack) < 9:
                num = 9 -len(stack)
                [stack.insert(len(stack)-1, '-') for n in range(num)]
                [classid.insert(len(classid)-1, '-') for n in range(num)]
            copy_cla = copy.deepcopy(classid)
            sample[ss] = copy_cla
            del stack[-1:]
            del classid[-1:]
        elif arr[3] == 'G':  
            handle('G', 8)
        elif arr[3] == 'F':
            handle('F', 7)
        elif arr[3] == 'O':
            handle('O', 6)
        elif arr[3] == 'C':
            handle('C', 5)
        elif arr[3] == 'P':
            handle('P', 4)
        elif arr[3] == 'D':
            handle('D', 3)

with open(outpre + '.species.xls', 'w') as out:
    out.write('U\tR\tD\tP\tC\tO\tF\tG\tS\tCount\n')
    if "Mycobacterium tuberculosis" not in sample.keys() and "Mycobacterium tuberculosis" in ks.keys():
        out.write(f'U_0_unclassified\tR_1_root\tD_2_Bacteria\tP_201174_Actinobacteria\tC_1760_Actinobacteria\tO_85007_Corynebacteriales\tF_1762_Mycobacteriaceae\tG_1763_Mycobacterium\tS_1773_Mycobacterium_tuberculosis\t{ks["Mycobacterium tuberculosis"]}\n')
    if "Mycobacterium avium" not in sample.keys() and "Mycobacterium avium" in ks.keys():
        out.write(f'U_0_unclassified\tR_1_root\tD_2_Bacteria\tP_201174_Actinobacteria\tC_1760_Actinobacteria\tO_85007_Corynebacteriales\tF_1762_Mycobacteriaceae\tG_1763_Mycobacterium\tS_1764_Mycobacterium_avium\t{ks["Mycobacterium avium"]}\n')
    for k,v in sample.items():
        if k in ks:
            out.write('\t'.join(v) + '\t' + str(ks[k]) + '\n')
out.close()

genus_count = defaultdict(int)
with open(outpre + '.species.xls') as file, open(outpre + '.genus.xls', 'w') as out:
    file.readline()
    out.write('U\tR\tD\tP\tC\tO\tF\tG\tCount\n')
    for line in file:
        i = line.strip().split('\t')
        k_str = "\t".join(i[0:8])
        genus_count[k_str] += int(i[9])
    for k,v in genus_count.items():
        out.write(f'{k}\t{v}\n')
